package com.redstar.interesting.common.util;

import com.redstar.interesting.common.constants.Constants;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;

/**
 * @author lihongxing
 */
public class ThreadMDCUtil {

    static final Logger log = LoggerFactory.getLogger(ThreadMDCUtil.class);

    public static void setTraceIdIfAbsent(){
        if (MDC.get(Constants.TRACE_ID) == null) {
            MDC.put(Constants.TRACE_ID, getTraceId());
        }
    }

    private static String getTraceId(){
        return UUID.randomUUID().toString().replaceAll("-", "");
    }

    public static void setTraceId(){
        MDC.put(Constants.TRACE_ID, getTraceId());
    }

    public static void setTraceId(String traceId) {
        MDC.put(Constants.TRACE_ID, traceId);
    }


    public static <T> Callable<T> wrap(final Callable<T> callable, final Map<String, String> context) {
        return ()->{
            if (context == null) {
                MDC.clear();
            }else {
                MDC.setContextMap(context);
            }
            setTraceIdIfAbsent();
            try {
                return callable.call();
            } catch (InterruptedException | ExecutionException exception) {
                log.error("ThreadPoolExecuteError", exception);
                return null;
            }finally {
                MDC.clear();
            }
        };
    }

    public static Runnable wrap(final Runnable runnable, final Map<String, String> context) {
        return ()->{
            if (context == null) {
                MDC.clear();
            }else {
                MDC.setContextMap(context);
            }
            setTraceIdIfAbsent();
            try {
                runnable.run();
            } catch (Exception exception) {
                log.error("ThreadPoolExecuteError", exception);
            }finally {
                MDC.clear();
            }
        };
    }

}
